import random
from typing import List, Tuple, Dict, Set, Any, Optional
import numpy as np

from env.grid_env import GridMap


class TaskGenerator:
    """
    任务生成器，用于生成机器人搬运任务
    """
    def __init__(self, grid_map: GridMap, config: Dict[str, Any]):
        """
        初始化任务生成器
        
        Args:
            grid_map: 网格地图
            config: 配置参数
        """
        self.grid_map = grid_map
        self.config = config
        
        # 从配置中获取取货点和送货点
        pickup_locations = config.get('pickup_locations', [])
        delivery_locations = config.get('delivery_locations', [])
        
        # 确保位置是元组而不是列表
        self.pickup_locations = [tuple(pos) for pos in pickup_locations] if pickup_locations else []
        self.delivery_locations = [tuple(pos) for pos in delivery_locations] if delivery_locations else []
        
        # 检测是否为仓储环境（通过检查障碍物分布特征）
        self.is_warehouse = self._detect_warehouse_environment()
        
        # 如果是仓储环境，生成货架旁边的位置作为取货点
        if self.is_warehouse:
            print("检测到仓储环境，生成货架旁边的取货点")
            self._generate_warehouse_pickup_locations()
        # 如果未提供位置，生成默认位置
        elif not self.pickup_locations:
            self._generate_default_pickup_locations()
        
        if not self.delivery_locations:
            self._generate_default_delivery_locations()
    
    def _detect_warehouse_environment(self) -> bool:
        """
        检测当前环境是否为仓储环境（通过障碍物分布特征）
        
        Returns:
            bool: 是否为仓储环境
        """
        # 检查障碍物分布是否符合仓储环境特征（规则排列的货架）
        obstacles = []
        for y in range(self.grid_map.height):
            for x in range(self.grid_map.width):
                if self.grid_map.grid[y, x] == 1:  # 障碍物
                    obstacles.append((x, y))
        
        # 如果障碍物数量在合理范围内，且分布规律（成行成列）
        if len(obstacles) > 20:  # 仓储环境通常有较多货架
            # 检查障碍物是否有规律排列
            x_coords = [x for x, y in obstacles]
            y_coords = [y for x, y in obstacles]
            
            # 计算x和y坐标的众数，检查是否有明显的行列分布
            x_counts = {}
            y_counts = {}
            for x in x_coords:
                x_counts[x] = x_counts.get(x, 0) + 1
            for y in y_coords:
                y_counts[y] = y_counts.get(y, 0) + 1
                
            # 如果有明显的行或列分布，判断为仓储环境
            max_x_count = max(x_counts.values()) if x_counts else 0
            max_y_count = max(y_counts.values()) if y_counts else 0
            
            return max_x_count > 3 or max_y_count > 3  # 如果某一行或列有3个以上障碍物
        
        return False
    
    def _generate_warehouse_pickup_locations(self, num_locations: int = 8):
        """
        在仓储环境中生成货架旁边的取货点
        
        Args:
            num_locations: 位置数量
        """
        self.pickup_locations = []
        obstacles = []
        
        # 找出所有障碍物（货架）
        for y in range(self.grid_map.height):
            for x in range(self.grid_map.width):
                if self.grid_map.grid[y, x] == 1:  # 障碍物
                    obstacles.append((x, y))
        
        # 随机选择一些货架，在其旁边生成取货点
        if obstacles:
            sampled_obstacles = random.sample(obstacles, min(len(obstacles), num_locations * 2))
            
            for obstacle_pos in sampled_obstacles:
                # 找到货架旁边的可通行位置
                for dx, dy in [(0, 1), (1, 0), (0, -1), (-1, 0)]:  # 上、右、下、左
                    nx, ny = obstacle_pos[0] + dx, obstacle_pos[1] + dy
                    pos = (nx, ny)
                    if self.grid_map.is_valid_position(pos) and pos not in self.pickup_locations:
                        self.pickup_locations.append(pos)
                        if len(self.pickup_locations) >= num_locations:
                            break
                
                if len(self.pickup_locations) >= num_locations:
                    break
        
        print(f"在仓储环境中生成了 {len(self.pickup_locations)} 个货架旁边的取货点")
    
    def _generate_default_pickup_locations(self, num_locations: int = 4):
        """
        生成默认取货点位置
        
        Args:
            num_locations: 位置数量
        """
        self.pickup_locations = []
        width, height = self.grid_map.width, self.grid_map.height
        
        # 在地图四个角附近生成取货点
        corners = [
            (width // 4, height // 4),
            (width // 4, 3 * height // 4),
            (3 * width // 4, height // 4),
            (3 * width // 4, 3 * height // 4)
        ]
        
        for x, y in corners:
            # 找到附近的有效位置
            for dx in range(-2, 3):
                for dy in range(-2, 3):
                    pos = (x + dx, y + dy)
                    if self.grid_map.is_valid_position(pos):
                        self.pickup_locations.append(pos)
                        break
                else:
                    continue
                break
        
        # 如果位置不足，随机生成
        while len(self.pickup_locations) < num_locations:
            x = random.randint(0, width - 1)
            y = random.randint(0, height - 1)
            pos = (x, y)
            if self.grid_map.is_valid_position(pos) and pos not in self.pickup_locations:
                self.pickup_locations.append(pos)
    
    def _generate_default_delivery_locations(self, num_locations: int = 4):
        """
        生成默认送货点位置
        
        Args:
            num_locations: 位置数量
        """
        self.delivery_locations = []
        width, height = self.grid_map.width, self.grid_map.height
        
        # 在地图中心和边缘生成送货点
        centers = [
            (width // 2, height // 2),
            (width // 2, 1),
            (1, height // 2),
            (width - 2, height - 2)
        ]
        
        for x, y in centers:
            # 找到附近的有效位置
            for dx in range(-2, 3):
                for dy in range(-2, 3):
                    pos = (x + dx, y + dy)
                    if self.grid_map.is_valid_position(pos):
                        self.delivery_locations.append(pos)
                        break
                else:
                    continue
                break
        
        # 如果位置不足，随机生成
        while len(self.delivery_locations) < num_locations:
            x = random.randint(0, width - 1)
            y = random.randint(0, height - 1)
            pos = (x, y)
            if self.grid_map.is_valid_position(pos) and pos not in self.delivery_locations:
                self.delivery_locations.append(pos)
    
    def generate_tasks(self, num_tasks: int, robot_positions: Dict[int, Tuple[int, int]]) -> Dict[int, Tuple[Tuple[int, int], Tuple[int, int]]]:
        """
        生成搬运任务
        
        Args:
            num_tasks: 任务数量
            robot_positions: 机器人初始位置 {robot_id: position}
            
        Returns:
            Dict[int, Tuple[Tuple[int, int], Tuple[int, int]]]: 任务字典 {robot_id: (start, goal)}
        """
        tasks = {}
        robot_ids = list(robot_positions.keys())
        
        # 确保任务数量不超过机器人数量
        num_tasks = min(num_tasks, len(robot_ids))
        
        # 随机选择机器人执行任务
        selected_robots = random.sample(robot_ids, num_tasks)
        
        for i, robot_id in enumerate(selected_robots):
            # 机器人起始位置
            start = robot_positions[robot_id]
            
            # 随机选择取货点和送货点
            pickup = random.choice(self.pickup_locations)
            delivery = random.choice(self.delivery_locations)
            
            # 确保取货点和送货点不同
            while pickup == delivery:
                delivery = random.choice(self.delivery_locations)
            
            # 分配任务
            tasks[robot_id] = (start, pickup)  # 第一阶段：从起点到取货点
        
        return tasks
    
    def generate_tasks_with_delivery(self, num_tasks: int, robot_positions: Dict[int, Tuple[int, int]]) -> Dict[int, List[Tuple[int, int]]]:
        """
        生成包含取货和送货的完整任务
        
        Args:
            num_tasks: 任务数量
            robot_positions: 机器人初始位置 {robot_id: position}
            
        Returns:
            Dict[int, List[Tuple[int, int]]]: 任务路径点列表 {robot_id: [start, pickup, delivery]}
        """
        tasks = {}
        robot_ids = list(robot_positions.keys())
        
        # 确保任务数量不超过机器人数量
        num_tasks = min(num_tasks, len(robot_ids))
        
        # 随机选择机器人执行任务
        selected_robots = random.sample(robot_ids, num_tasks)
        
        for i, robot_id in enumerate(selected_robots):
            # 机器人起始位置
            start = robot_positions[robot_id]
            
            # 随机选择取货点和送货点
            pickup = random.choice(self.pickup_locations)
            delivery = random.choice(self.delivery_locations)
            
            # 确保取货点和送货点不同
            while pickup == delivery:
                delivery = random.choice(self.delivery_locations)
            
            # 分配任务
            tasks[robot_id] = [start, pickup, delivery]  # 完整路径：起点 -> 取货点 -> 送货点
        
        return tasks
    
    def generate_balanced_tasks(self, num_tasks: int, robot_positions: Dict[int, Tuple[int, int]]) -> Dict[int, Tuple[Tuple[int, int], Tuple[int, int]]]:
        """
        生成平衡的搬运任务，尽量使每个机器人的任务难度相似
        
        Args:
            num_tasks: 任务数量
            robot_positions: 机器人初始位置 {robot_id: position}
            
        Returns:
            Dict[int, Tuple[Tuple[int, int], Tuple[int, int]]]: 任务字典 {robot_id: (start, goal)}
        """
        tasks = {}
        robot_ids = list(robot_positions.keys())
        
        # 确保任务数量不超过机器人数量
        num_tasks = min(num_tasks, len(robot_ids))
        
        # 确保pickup_locations中的元素是元组而不是列表
        self.pickup_locations = [tuple(pos) if isinstance(pos, list) else pos for pos in self.pickup_locations]
        
        # 计算每个机器人到每个取货点的距离
        distances = {}
        for robot_id, start in robot_positions.items():
            distances[robot_id] = []
            for pickup in self.pickup_locations:
                # 计算曼哈顿距离
                dist = abs(start[0] - pickup[0]) + abs(start[1] - pickup[1])
                distances[robot_id].append((dist, pickup))
            
            # 按距离排序
            distances[robot_id].sort()
        
        # 分配任务，优先分配距离适中的任务
        selected_robots = random.sample(robot_ids, num_tasks)
        assigned_pickups = set()
        
        for robot_id in selected_robots:
            start = robot_positions[robot_id]
            
            # 选择距离适中的取货点（避免最近和最远的）
            pickup_options = distances[robot_id]
            middle_index = len(pickup_options) // 2
            
            # 尝试分配未被分配的取货点
            pickup = None
            for i in range(len(pickup_options)):
                idx = (middle_index + i) % len(pickup_options)
                _, candidate = pickup_options[idx]
                if candidate not in assigned_pickups:
                    pickup = candidate
                    assigned_pickups.add(pickup)
                    break
            
            # 如果所有取货点都已分配，随机选择一个
            if pickup is None:
                pickup = random.choice(self.pickup_locations)
            
            # 分配任务
            tasks[robot_id] = (start, pickup)
        
        return tasks 